package com.baishancloud.log.common.connector.es

import com.baishancloud.log.common.util.ParamUtil.paramName
import org.apache.flink.api.common.functions.RuntimeContext
import org.apache.flink.api.java.utils.ParameterTool
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkBase.FlushBackoffType
import org.apache.flink.streaming.connectors.elasticsearch.util.RetryRejectedExecutionFailureHandler
import org.apache.flink.streaming.connectors.elasticsearch.{ElasticsearchSinkFunction, RequestIndexer}
import org.apache.flink.streaming.connectors.elasticsearch7.{ElasticsearchSink, RestClientFactory}
import org.apache.http.HttpHost
import org.apache.http.auth.{AuthScope, UsernamePasswordCredentials}
import org.apache.http.client.config.RequestConfig
import org.apache.http.impl.client.BasicCredentialsProvider
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder
import org.elasticsearch.action.index.IndexRequest
import org.elasticsearch.client.{Requests, RestClientBuilder}
import org.elasticsearch.common.xcontent.XContentType
import org.slf4j.{Logger, LoggerFactory}

import java.util
import scala.beans.BeanProperty


/**
 *
 *
 * @author ziqiang.wang,klein
 * @date 2021/10/26 19:16
 */
class ESSinkConnector7 extends Serializable {


  private val LOGGER: Logger = LoggerFactory.getLogger(classOf[ESSinkConnector7])

  /**
   * 必选<br>
   * 连接es的集群的一个或多个主机地址，格式：'http://host_name:9092;http://host_name:9093'
   */
  @BeanProperty var httpHosts: util.List[HttpHost] = new util.ArrayList[HttpHost]()

  def httpHosts_=(urls: String): Unit = {
    urls.split(";").foreach(url => {
      val schemaHostPort: Array[String] = url.trim.split("://")
      val hostPort: Array[String] = schemaHostPort(1).trim.split(":")
      if (hostPort.length == 2) {
        httpHosts.add(new HttpHost(hostPort(0), hostPort(1).toInt, schemaHostPort(0)))
      } else {
        httpHosts.add(new HttpHost(hostPort(0)))
      }
    })
  }

  /**
   * 必选<br>
   * 索引前缀<br>
   * 比如："myIndex-"；
   */
  @BeanProperty var index: String = _
  /**
   * 写入ES之前最多缓存多少条数据，默认1000条
   */
  @BeanProperty var bulkFlushMaxActions: Int = 1000

  /**
   * 写入ES之前最多缓存多少size的数据，默认2M，单位：M
   */
  @BeanProperty var bulkFlushMaxSizeMb: Int = 2
  /**
   * 写入ES的时间间隔，不管是否达到了数量或size的条件，到指定时间间隔都会触发flush，默认为1秒，单位：秒
   */
  @BeanProperty var bulkFlushInterval: Long = 1000
  /**
   * 是否在checkpoint成功之后将数据flush到es，默认为false。如果设置了文档id，在多次写入的情况下，不会有数据重复。如果没有文档id，则可能造成写入数据重复，开启checkpoint可避免重复写入。
   */
  @BeanProperty var flushOnCheckpoint: Boolean = false
  /**
   * 写入ES错误的重试次数,默认8次
   */
  @BeanProperty var bulkFlushBackoffRetries: Int = 8
  /**
   * ES7认证用户 不为空则开启
   */
  @BeanProperty var username: String = ""
  /**
   * ES7认证密码 不为空则开启
   */
  @BeanProperty var password: String = ""


  /**
   * 返回一个ESSink，接收二元组(indexSuffix,data)，动态索引
   *
   * @return ElasticsearchSink[(String,String)]
   */
  def dynamicIndex(): ElasticsearchSink[(String, String)] = {
    val builder = new ElasticsearchSink.Builder[(String, String)](httpHosts,
      new ElasticsearchSinkFunction[(String, String)]() {
        override def process(element: (String, String), ctx: RuntimeContext, indexer: RequestIndexer): Unit = {
          val request: IndexRequest = Requests.indexRequest
            .index(s"${index}${element._1}")
            .source(element._2, XContentType.JSON)
          indexer.add(request)
        }
      }


    )
    otherParam(builder)
    val esSink: ElasticsearchSink[(String, String)] = builder.build()
    if (!flushOnCheckpoint) {
      esSink.disableFlushOnCheckpoint()
    }
    esSink
  }

  /**
   * 返回一个ESSink，接收三元组(id,indexSuffix,data)，动态索引
   *
   * @return ElasticsearchSink[(String,String,String)]
   */
  def dynamicIndexHaveId(): ElasticsearchSink[(String, String, String)] = {
    val builder = new ElasticsearchSink.Builder[(String, String, String)](httpHosts,
      new ElasticsearchSinkFunction[(String, String, String)]() {
        override def process(element: (String, String, String), ctx: RuntimeContext, indexer: RequestIndexer): Unit = {
          val request: IndexRequest = Requests.indexRequest
            .id(element._1)
            .index(s"${index}${element._2}")
            .source(element._3, XContentType.JSON)
          indexer.add(request)
        }
      }


    )
    otherParam(builder)
    val esSink: ElasticsearchSink[(String, String, String)] = builder.build()
    if (!flushOnCheckpoint) {
      esSink.disableFlushOnCheckpoint()
    }
    esSink
  }

  /**
   * 其他一些参数的调整，主要涉及到性能等一些优化
   */
  private def otherParam(builder: ElasticsearchSink.Builder[_]): Unit = {
    builder.setRestClientFactory(
      new RestClientFactory() {
        override def configureRestClientBuilder(restClientBuilder: RestClientBuilder): Unit = {
          if (username != "" && password != "") {
            //用户名、密码
            restClientBuilder.setHttpClientConfigCallback(new RestClientBuilder.HttpClientConfigCallback() {
              override def customizeHttpClient(httpClientBuilder: HttpAsyncClientBuilder): HttpAsyncClientBuilder = {
                val credentialsProvider = new BasicCredentialsProvider
                credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(username, password))
                httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider)
              }
            })
          }
          //一些超时时间设置
          restClientBuilder.setRequestConfigCallback(new RestClientBuilder.RequestConfigCallback() {
            override def customizeRequestConfig(requestConfigBuilder: RequestConfig.Builder): RequestConfig.Builder = {
              requestConfigBuilder.setConnectTimeout(50000)
              requestConfigBuilder.setSocketTimeout(300000)
              requestConfigBuilder.setConnectionRequestTimeout(10000)
              requestConfigBuilder
            }
          })
        }
      }

    )
    builder.setBulkFlushMaxActions(bulkFlushMaxActions)
    builder.setBulkFlushMaxSizeMb(bulkFlushMaxSizeMb)
    builder.setBulkFlushInterval(bulkFlushInterval)
    builder.setBulkFlushBackoff(true)
    builder.setBulkFlushBackoffDelay(1000)
    builder.setBulkFlushBackoffRetries(bulkFlushBackoffRetries)
    builder.setBulkFlushBackoffType(FlushBackoffType.CONSTANT)
    builder.setFailureHandler(new RetryRejectedExecutionFailureHandler())
  }

}


object ESSinkConnector7 extends Serializable {

  /**
   * @param number 要获取那个数字对应的数据源
   */
  def builder(parameterTool: ParameterTool, number: String = ""): ESSinkConnector7 = {
    val connector: ESSinkConnector7 = new ESSinkConnector7()

    if (parameterTool.has(paramName(sinkHostsP, number))) {
      connector.httpHosts = parameterTool.get(paramName(sinkHostsP, number))
    } else {
      throw new RuntimeException(s"未设置 --${paramName(sinkHostsP, number)} 参数")
    }

    if (parameterTool.has(paramName(sinkIndexP, number))) {
      connector.setIndex(parameterTool.get(paramName(sinkIndexP, number)))
    } else {
      throw new RuntimeException(s"未设置 --${paramName(sinkIndexP, number)} 参数")
    }

    if (parameterTool.has(paramName(sinkBulkFlushMaxActionsP, number))) {
      connector.setBulkFlushMaxActions(parameterTool.getInt(paramName(sinkBulkFlushMaxActionsP, number)))
    }
    if (parameterTool.has(paramName(sinkBulkFlushMaxSizeMbP, number))) {
      connector.setBulkFlushMaxSizeMb(parameterTool.getInt(paramName(sinkBulkFlushMaxSizeMbP, number)))
    }
    if (parameterTool.has(paramName(sinkBulkFlushIntervalP, number))) {
      connector.setBulkFlushInterval(parameterTool.getLong(paramName(sinkBulkFlushIntervalP, number) * 1000))
    }
    if (parameterTool.has(paramName(sinkFlushOnCheckpointP, number))) {
      connector.setFlushOnCheckpoint(parameterTool.getBoolean(paramName(sinkFlushOnCheckpointP, number)))
    }
    if (parameterTool.has(paramName(sinkBulkFlushBackoffRetriesP, number))) {
      connector.setBulkFlushBackoffRetries(parameterTool.getInt(paramName(sinkBulkFlushBackoffRetriesP, number)))
    }
    if (parameterTool.has(paramName(sinkUsernameP, number))) {
      connector.setUsername(parameterTool.get(paramName(sinkUsernameP, number)))
    }
    if (parameterTool.has(paramName(sinkPasswordP, number))) {
      connector.setPassword(parameterTool.get(paramName(sinkPasswordP, number)))
    }

    connector
  }
}


